#!/bin/bash

PORT=7001
BASE_DIR="./received_files"
LOG_FILE="./file_transfer.log"
TEMP_DIR="$BASE_DIR/temp"

# 初始化目录
mkdir -p "$BASE_DIR" "$TEMP_DIR"
touch "$LOG_FILE"

log_message() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] $message" | tee -a "$LOG_FILE"
}

# 检查端口是否可用
check_port() {
    if netstat -tuln 2>/dev/null | grep -q ":$PORT "; then
        log_message "错误: 端口 $PORT 已被占用"
        return 1
    fi
    return 0
}

# 检查必要工具
check_dependencies() {
    local deps=("file" "nc" "unzip")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            log_message "错误: 未找到 $dep 命令"
            return 1
        fi
    done
    return 0
}

# 检测文件类型
detect_file_type() {
    local file="$1"
    local file_output=$(file -b "$file" 2>/dev/null)
    
    case "$file_output" in
        *"Zip archive"*)
            echo "zip"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

# 清理历史文件
cleanup_old_files() {
    log_message "清理历史文件..."
    rm -rf "$BASE_DIR"/*.zip
    rm -rf "$TEMP_DIR"/*
    log_message "历史文件清理完成"
}

# 处理 ZIP 文件
process_zip() {
    local file="$1"
    local filename="$2"
    
    log_message "开始处理 ZIP 文件: $filename"
    
    # 解压到 /root/ 目录
    if unzip -o "$file" -d "/root/" > /dev/null 2>&1; then
        log_message "ZIP 文件解压成功: $filename -> /root/"
        
        # 删除 ZIP 文件
        rm -f "$file"
        log_message "已删除 ZIP 文件: $filename"
        
        # 检查并运行 WEI-iEdge_update.sh 脚本
        if [ -f "/root/WEI-iEdge_update.sh" ]; then
            log_message "找到 WEI-iEdge_update.sh 脚本，正在后台执行..."
            # 后台运行 WEI-iEdge_update.sh 脚本
            nohup /bin/bash /root/WEI-iEdge_update.sh > /root/WEI-iEdge/build/logs/ota_execution.log 2>&1 &
            log_message "WEI-iEdge_update.sh 脚本已在后台启动，日志输出到 /root/WEI-iEdge/build/logs/ota_execution.log"
        else
            log_message "警告: 在 /root/ 目录下未找到 WEI-iEdge_update.sh 脚本"
        fi
        
        return 0
    else
        log_message "错误: ZIP 文件解压失败: $filename"
        return 1
    fi
}

# 主接收循环
receive_loop() {
    while true; do
        log_message "等待客户端连接..."
        
        # 清理历史文件
        cleanup_old_files
        
        # 使用临时文件接收数据
        TEMP_FILE=$(mktemp -p "$TEMP_DIR")
        
        # 监听端口并接收数据
        if nc -l -p $PORT > "$TEMP_FILE" 2>/dev/null; then
            if [ -s "$TEMP_FILE" ]; then
                local file_size=$(stat -c%s "$TEMP_FILE")
                log_message "接收到数据流 (大小: $file_size 字节)"
                
                # 检测文件类型
                local file_type=$(detect_file_type "$TEMP_FILE")
                
                if [ "$file_type" = "zip" ]; then
                    # 生成带时间戳的文件名
                    local timestamp=$(date +%Y%m%d_%H%M%S)
                    local final_filename="update_${timestamp}.zip"
                    local final_file="$BASE_DIR/$final_filename"
                    
                    mv "$TEMP_FILE" "$final_file"
                    
                    log_message "ZIP 文件保存为: $final_filename"
                    
                    # 处理 ZIP 文件
                    process_zip "$final_file" "$final_filename"
                else
                    log_message "错误: 接收到非 ZIP 格式文件，已删除"
                    rm -f "$TEMP_FILE"
                fi
            else
                rm -f "$TEMP_FILE"
                log_message "接收到空文件"
            fi
        else
            log_message "接收连接失败"
            sleep 1
        fi
    done
}

# 显示统计信息
show_stats() {
    log_message "=== 文件接收统计 ==="
    log_message "总接收文件数: $(find "$BASE_DIR" -maxdepth 1 -type f -name "*.zip" | wc -l)"
}

# 主函数
main() {
    log_message "=== 启动 ZIP 文件接收服务 ==="
    log_message "监听端口: $PORT"
    log_message "文件保存目录: $BASE_DIR"
    log_message "解压目标: /root/"
    
    if ! check_port; then
        exit 1
    fi
    
    if ! check_dependencies; then
        exit 1
    fi
    
    # 显示初始统计
    show_stats
    
    # 设置信号处理
    trap 'log_message "服务停止"; show_stats; exit 0' INT TERM
    
    # 启动时清理历史文件
    cleanup_old_files
    
    # 启动接收循环
    receive_loop
}

# 运行主函数
main

